# -*- coding: utf-8 -*-
# @Author: bao
# @Date:   2017-03-31 15:10:31
# @Last Modified by:   bao
# @Last Modified time: 2017-05-12 09:56:13
import sys
sys.path.append("..")

from datetime import datetime
import re
import random
import copy
from CoreSim.MScheduler import *
from CoreSim.Flow import *
from CoreSim.Event import *
from CoreSim.Link import *
from CoreSim.Unit import *

def getSortedFlows(activeFlows, priority = 'flowSize', reverse=False):
    if priority == 'remainTime':
        sortedActiveFlows = sorted(activeFlows, key=lambda flow: flow.remainTime, reverse=reverse)
    elif priority == 'remainSize':
        sortedActiveFlows = sorted(activeFlows, key=lambda flow: flow.remainSize, reverse=reverse)
    elif priority == 'flowSize':
        sortedActiveFlows = sorted(activeFlows, key=lambda flow: flow.flowSize, reverse=reverse)

    return sortedActiveFlows

def assignNewMulticastFlow(srcHost, dstHosts, flows, flowSize, startTime, mcastFlowId, numHosts, numToRs, numCores):
    # multicast flow
    flow = Flow()
    flow.startId = srcHost
    flow.endId = dstHosts
    flow.SetFlowSize(flowSize)
    flow.startTime = startTime
    flow.flowId = len(flows) + 1
    flow.mcastFlowId = mcastFlowId
    flows.append(flow)

    return flow

class MSchedulerJcast(MScheduler):
    '''
    Multicast scheduler
    Constructing multicast tree
    '''
    def __init__(self, topo_args, flow_args, conf_args):
        MScheduler.__init__(self, topo_args, flow_args, conf_args)
        self.type = 'Jcast'

        self.runningMcastFlows = []
        self.mcastToR = {}              # ToRId: {engress portId: mcastflowId}

    def InitialMulticast(self, eventQueue):
        for mflowId in self.mcastFlows:
            mcastFlow = self.mcastFlows[mflowId]
            srcHost = mcastFlow.srcHost
            dstHosts = mcastFlow.dstHosts.keys()
            
            flow = assignNewMulticastFlow(srcHost, dstHosts, self.flows, mcastFlow.flowSize, mcastFlow.startTime, \
                mcastFlow.flowId, self.numHosts, self.numToRs, self.numCores)

            # eventqueue
            event = FlowArrivalEvent(flow.startTime, flow)
            eventQueue.push(event, flow.startTime)

    def MulticastFlowArrival(self, curFlow, curTime, eventQueue):
        #print 'MulticastFlowArrival'
        if self.UpdateTopologyEvent:
            self.UpdateTopologyEvent.cancelled = True
        if self.LinkAdjustedEvent:
            self.LinkAdjustedEvent.cancelled = True

        starttime = datetime.now()
        # update the routing and links
        for flow in self.runningMcastFlows:
            pathInLink = flow.pathLinkIds
            for link_id in pathInLink:
                link = self.Links[link_id]
                link.flowIds.remove(flow.flowId)

        # reassign ocs links
        self.mcastToR = {}
        self.ucastToR = copy.deepcopy(self.ucastToRFully)
        self.ucastPort = copy.deepcopy(self.ucastPortFully)

        self.runningMcastFlows.append(curFlow)
        sortedRunningFlows = getSortedFlows(self.runningMcastFlows, priority='remainSize', reverse=True)
        for flow in sortedRunningFlows:
            mcastFlow = self.mcastFlows[flow.mcastFlowId]
            srcHost = mcastFlow.srcHost
            srcToR = mcastFlow.srcToR
            dstToRs = mcastFlow.dstToRs.keys()
            #print '----------'
            #print srcToR, dstToRs

            free_ports = {}
            free_tors = []
            for portId in self.ucastToR[srcToR][0]:
                free_ports[portId] = self.ucastPort[portId][1]
                free_tors = list(set(free_tors)| set(self.ucastPort[portId][1]))

            flow.pathLinkIds = []
            flow.ocsPathLinkIds = []
            srcToRNodeId = ConvertToNodeId(srcToR, ToR, self.numHosts*self.numToRs, self.numToRs)
            
            if set(dstToRs).issubset(free_tors): # assign ocs links 

                flow.ocsPathLinkIds.append((srcHost, srcToRNodeId))

                free_portIds = sorted(free_ports, key=lambda x : len(free_ports[x])) #
                dstToRs_Set = set(dstToRs)

                for portId in free_portIds:
                    torIds = list(dstToRs_Set & set(free_ports[portId]))
                    dstToRs_Set -= set(free_ports[portId])
                    #print torIds
                    
                    self.ucastToR[srcToR][0].remove(portId)
                    self.ucastPort[portId][0].remove(srcToR)
                    self.mcastToR[srcToR] = {}
                    self.mcastToR[srcToR][portId] = flow.mcastFlowId
                        
                    ocsNodeId = ConvertToNodeId(portId, OCS, self.numHosts*self.numToRs, self.numToRs, self.numCores)
                    flow.ocsPathLinkIds.append((srcToRNodeId, ocsNodeId))
                    for torId in torIds:
                        self.ucastToR[torId][1].remove(portId)
                        self.ucastPort[portId][1].remove(torId)
                        dstToRNodeId = ConvertToNodeId(torId, ToR, self.numHosts*self.numToRs, self.numToRs)
                        flow.ocsPathLinkIds.append((ocsNodeId, dstToRNodeId))

                        dstHosts = mcastFlow.dstToRs[torId]
                        for dstHost in dstHosts:
                            flow.ocsPathLinkIds.append((dstToRNodeId, dstHost))
                    
                    #print flow.ocsPathLinkIds
                    if not dstToRs_Set:
                        break
                #print self.ucastToR
                #print self.ucastPort

            # eps routes (if has ocs routes, backup for reconfiguration)
            flow.pathLength = 4
            coreId = random.randint(1, self.numCores)
            coreNodeId = ConvertToNodeId(coreId, Core, self.numHosts*self.numToRs, self.numToRs)
            flow.pathLinkIds.append((srcHost, srcToRNodeId))
            flow.pathLinkIds.append((srcToRNodeId, coreNodeId))

            for torId in dstToRs:
                dstToRNodeId = ConvertToNodeId(torId, ToR, self.numHosts*self.numToRs, self.numToRs)
                flow.pathLinkIds.append((coreNodeId, dstToRNodeId))
                
                dstHosts = mcastFlow.dstToRs[torId]
                for dstHost in dstHosts:
                    flow.pathLinkIds.append((dstToRNodeId, dstHost))

            pathInLink = flow.pathLinkIds
            for link_id in pathInLink:
                link = self.Links[link_id]
                if flow.flowId not in link.flowIds:
                    link.flowIds.append(flow.flowId)

        endtime = datetime.now()
        timediff = endtime - starttime
        print 'jcast time: ', timediff
        # update the unicast/multicast links and routes
        self.UpdateTopology(eventQueue, curTime, 'multicast')

    def UpdateMulticast(self, mflowId, curFlow, curTime, eventQueue):
        mcastFlow = self.mcastFlows[mflowId]
        mcastFlow.finishTime = curTime
        mcastFlow.energy += curFlow.energy
        print mcastFlow.energy
        print 'mcastFlow:%s transmission time %s energy %s'%(mflowId, mcastFlow.finishTime-mcastFlow.startTime, mcastFlow.energy)
        
        #print 'MulticastFlowFinished'
        self.runningMcastFlows.remove(curFlow) 
        if self.UpdateTopologyEvent:
            self.UpdateTopologyEvent.cancelled = True
        if self.LinkAdjustedEvent:
            self.LinkAdjustedEvent.cancelled = True

        # update the routing and links
        for flow in self.runningMcastFlows:
            pathInLink = flow.pathLinkIds
            for link_id in pathInLink:
                link = self.Links[link_id]
                link.flowIds.remove(flow.flowId)

        # reassign ocs links
        self.mcastToR = {}
        self.ucastToR = copy.deepcopy(self.ucastToRFully)
        self.ucastPort = copy.deepcopy(self.ucastPortFully)

        sortedRunningFlows = getSortedFlows(self.runningMcastFlows, priority='remainSize', reverse=True)
        for flow in sortedRunningFlows:
            mcastFlow = self.mcastFlows[flow.mcastFlowId]
            srcHost = mcastFlow.srcHost
            srcToR = mcastFlow.srcToR
            dstToRs = mcastFlow.dstToRs.keys()
            #print '----------'
            #print srcToR, dstToRs

            free_ports = {}
            free_tors = []
            for portId in self.ucastToR[srcToR][0]:
                free_ports[portId] = self.ucastPort[portId][1]
                free_tors = list(set(free_tors)| set(self.ucastPort[portId][1]))

            flow.pathLinkIds = []
            flow.ocsPathLinkIds = []
            srcToRNodeId = ConvertToNodeId(srcToR, ToR, self.numHosts*self.numToRs, self.numToRs)
            
            if set(dstToRs).issubset(free_tors): # assign ocs links 

                flow.ocsPathLinkIds.append((srcHost, srcToRNodeId))

                free_portIds = sorted(free_ports, key=lambda x : len(free_ports[x])) #
                dstToRs_Set = set(dstToRs)

                for portId in free_portIds:
                    torIds = list(dstToRs_Set & set(free_ports[portId]))
                    dstToRs_Set -= set(free_ports[portId])
                    #print torIds
                    
                    self.ucastToR[srcToR][0].remove(portId)
                    self.ucastPort[portId][0].remove(srcToR)
                    self.mcastToR[srcToR] = {}
                    self.mcastToR[srcToR][portId] = flow.mcastFlowId
                        
                    ocsNodeId = ConvertToNodeId(portId, OCS, self.numHosts*self.numToRs, self.numToRs, self.numCores)
                    flow.ocsPathLinkIds.append((srcToRNodeId, ocsNodeId))
                    for torId in torIds:
                        self.ucastToR[torId][1].remove(portId)
                        self.ucastPort[portId][1].remove(torId)
                        dstToRNodeId = ConvertToNodeId(torId, ToR, self.numHosts*self.numToRs, self.numToRs)
                        flow.ocsPathLinkIds.append((ocsNodeId, dstToRNodeId))

                        dstHosts = mcastFlow.dstToRs[torId]
                        for dstHost in dstHosts:
                            flow.ocsPathLinkIds.append((dstToRNodeId, dstHost))
                    
                    #print flow.ocsPathLinkIds
                    if not dstToRs_Set:
                        break
                #print self.ucastToR
                #print self.ucastPort

            # eps routes (if has ocs routes, backup for reconfiguration)
            coreId = random.randint(1, self.numCores)
            coreNodeId = ConvertToNodeId(coreId, Core, self.numHosts*self.numToRs, self.numToRs)
            flow.pathLinkIds.append((srcHost, srcToRNodeId))
            flow.pathLinkIds.append((srcToRNodeId, coreNodeId))
            
            for torId in dstToRs:
                dstToRNodeId = ConvertToNodeId(torId, ToR, self.numHosts*self.numToRs, self.numToRs)
                flow.pathLinkIds.append((coreNodeId, dstToRNodeId))
                
                dstHosts = mcastFlow.dstToRs[torId]
                for dstHost in dstHosts:
                    flow.pathLinkIds.append((dstToRNodeId, dstHost))

            pathInLink = flow.pathLinkIds
            for link_id in pathInLink:
                link = self.Links[link_id]
                if flow.flowId not in link.flowIds:
                    link.flowIds.append(flow.flowId)

        # update the unicast/multicast links and routes
        self.UpdateTopology(eventQueue, curTime, 'multicast')
        

